import tensorrt as trt
import os
import onnx

class TRTExporter:
    def __init__(self, onnx_path):
        self.onnx_path = onnx_path
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.builder = trt.Builder(self.logger)
        self.network = self.builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
        self.config = self.builder.create_builder_config()

        self.config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 28)  # 256MB
        self.config.set_flag(trt.BuilderFlag.FP16)  # 可选：注释掉以使用 FP32

    def build_engine(self, engine_path):
        # 验证 ONNX 模型
        model = onnx.load(self.onnx_path)
        onnx.checker.check_model(model)
        print("ONNX model is valid.")

        # 解析 ONNX 模型
        parser = trt.OnnxParser(self.network, self.logger)
        with open(self.onnx_path, 'rb') as model:
            if not parser.parse(model.read()):
                print("ERROR: Failed to parse the ONNX model.")
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return False
            else:
                print("ONNX model parsed successfully.")
                for error in range(parser.num_errors):
                    print("Warning:", parser.get_error(error))

        # 检查网络输入输出
        print("Number of inputs:", self.network.num_inputs)
        print("Number of outputs:", self.network.num_outputs)
        if self.network.num_inputs == 0:
            print("ERROR: Network has no inputs. Parsing may have failed.")
            return False

        # 设置优化配置文件
        if self.network.num_inputs > 0:
            input_name = self.network.get_input(0).name
            print("Input name:", input_name)
            profile = self.builder.create_optimization_profile()
            min_shape = (1, 3, 640, 640)  # 可调整为 (1, 3, 320, 320)
            opt_shape = (1, 3, 640, 640)
            max_shape = (1, 3, 640, 640)  # 减小 batch size，例如 (1, 3, 640, 640)
            profile.set_shape(input_name, min_shape, opt_shape, max_shape)
            self.config.add_optimization_profile(profile)
            print("Optimization profile added successfully.")
        else:
            print("ERROR: No inputs found in the network.")
            return False

        # 构建引擎
        plan = self.builder.build_serialized_network(self.network, self.config)
        if plan is None:
            print("ERROR: Failed to build serialized network.")
            return False
        with open(engine_path, 'wb') as f:
            f.write(plan)
        return True

if __name__ == '__main__':
    exporter = TRTExporter('models/best.onnx')
    if exporter.build_engine('models/yolov8n.trt'):
        print('TensorRT build success.')
    else:
        print('Failed to build.')
